import torch.nn as nn
import efficet
import docrnn
import fitter
import bfb
import torch
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.eff=efficet.get_efficientnet_v2('efficientnet_v2_s')
        self.doc=docrnn.DOCRNN()
        self.bfb=bfb.bfb()
    def forward(self, x):
        b = fitter.fitterload(x,self.training)
        x1 = self.eff(torch.unsqueeze(torch.unsqueeze(torch.from_numpy(b[0]).cuda(), dim=0), dim=0))
        x2 = self.doc(torch.unsqueeze(torch.unsqueeze(torch.from_numpy(b[1]).cuda(), dim=0), dim=0))
        x = self.bfb(x1,x2)
        return x
